from dotenv import load_dotenv
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import (
    ChatPromptTemplate,
)
#执行器引用
from langchain_experimental.utilities import PythonREPL
from langchain_openai import ChatOpenAI

load_dotenv()

template= """Write some python code to solve the user's problem.
Return only python code in Markdown format,  e.g.:
``` python
...
```"""
prompt = ChatPromptTemplate.from_messages([("system", template),("human", "{input}")])

model = ChatOpenAI()
#输出处理
def _sanitize_output(text: str):
     _,after = text.split("'''python")
     return after.split("```")[0]
chain = prompt |model | StrOutputParser() | _sanitize_output | PythonREPL().run
#chain = prompt |model | StrOutputParser()
result = chain.invoke({"input":"Whats 2 plus 2"})

print(result)




